"""tensor的变形操作"""

import torch


def test_reshape():
    print("=========================== reshape ====")
    a = torch.rand(2, 3)
    print("a: ")
    print(a)
    print(a.shape)
    print()

    print("torch.reshape(a, (3,2)): ")
    out = torch.reshape(a, (3, 2))
    print(out)
    print(out.shape)
    print()


def test_t_transpose():
    print("===================== t ====")
    a = torch.rand(2, 3)
    print("a: ")
    print(a)
    print(a.shape)
    print()

    print("torch.t(out): 转置 ")
    t = torch.t(a)
    print(t)
    print(t.shape)
    print()

    print("=========================== transpose ====")
    print("torch.transpose(out, 0, 1): 交换维度0和1 对二维来说，等同于转置 ")
    trans = torch.transpose(a, 0, 1)
    print(trans)
    print(trans.shape)
    print()

    print("更高维度(1,2,3)的tensor进行transpose ---> ")
    print("torch.rand(1,2,3)")
    a = torch.rand(1, 2, 3)
    print(a)
    print(a.shape)
    print()

    print("torch.transpose(a, 0, 1)  a.shape=(1,2,3)")
    out = torch.transpose(a, 0, 1)
    print(out)
    print(out.shape)
    print()


def test_squeeze():
    print("================================ squeeze ====")
    print("torch.rand(2,1,3)")
    a = torch.rand(2, 1, 3)
    print(a)
    print(a.shape)
    print()

    print("==== squeeze 将维度=1的维度删除 ====")
    print("torch.squeeze(a) a.shape=(2,1,3), squeeze后shape=(2,3) ")
    out = torch.squeeze(a)
    print(out)
    print(out.shape)
    print()

    print("==== unsqueeze 扩展维度 ====")
    print("torch.unsqueeze(out, -1) out.shape=(2,3), -1表示在最后一个维度进行拓展，unsqueeze后shape=(2,3, 1) ")
    out = torch.unsqueeze(out, -1)
    print(out)
    print(out.shape)
    print()


def test_unbind():
    print("================================= unbind ====")
    print("torch.rand(1,2,3)")
    a = torch.rand(1, 2, 3)
    print(a)
    print(a.shape)
    print()

    print("torch.unbind(a, 1) 按照dim=1将tensor拆分")
    out = torch.unbind(a, 1)
    print(out)
    print()


def test_flip_rot90():
    print("================================= flip ====")
    print("torch.rand(1,2,3)")
    a = torch.rand(1, 2, 3)
    print(a)
    print(a.shape)
    print()

    print("torch.flip(a, [2])   对维度2进行翻转")
    out = torch.flip(a, [2])
    print(out)
    print()

    print("torch.flip(a, [1, 2])   先对维度1,再对维度2进行翻转")
    out = torch.flip(a, [1, 2])
    print(out)
    print()

    print("================================== rot90 ====")
    print("torch.rot90(a, 1, dims=[0, 2]), k=几，就是旋转几个90度，正值逆时针，负值顺时针")
    out = torch.rot90(a, 1, dims=[0, 2])
    print(out)
    print(out.shape)


if __name__ == '__main__':
    test_reshape()
    test_t_transpose()
    test_squeeze()
    test_unbind()
    test_flip_rot90()
